﻿' SearchEngine.vb
Imports System.Text.RegularExpressions
Imports TodoList.Models

Public Class SearchEngine
    Private ReadOnly _indexBuilder As IndexBuilder

    Public Sub New(indexBuilder As IndexBuilder)
        _indexBuilder = indexBuilder
    End Sub

    ' 执行搜索
    ' 修改 Search 方法以突出显示图片 OCR 结果
    Public Function Search(query As String, Optional fileType As String = "", Optional modifiedAfter As DateTime? = Nothing, Optional includeImages As Boolean = True) As List(Of SearchResult)
        ' 获取索引
        Dim index = _indexBuilder.GetIndex()
        If index Is Nothing OrElse index.Count = 0 Then
            Return New List(Of SearchResult)()
        End If

        ' 预处理查询
        Dim queryTerms = PreprocessQuery(query)
        If queryTerms.Count = 0 Then
            Return New List(Of SearchResult)()
        End If

        ' 存储搜索结果
        Dim results As New List(Of SearchResult)()

        ' 对每个文件执行搜索
        For Each fileInfo In index.Values
            ' 应用过滤器
            If Not String.IsNullOrEmpty(fileType) AndAlso Not fileInfo.Extension.Equals(fileType, StringComparison.OrdinalIgnoreCase) Then
                Continue For
            End If

            If modifiedAfter.HasValue AndAlso fileInfo.Modified < modifiedAfter.Value Then
                Continue For
            End If

            ' 如果不包括图片且当前是图片文件，则跳过
            If Not includeImages AndAlso IsImageFile(fileInfo.Extension) Then
                Continue For
            End If

            ' 计算得分
            Dim score As Double = 0
            Dim matchedLines As New List(Of String)()

            ' 文件名匹配
            For Each term In queryTerms
                If fileInfo.Name.IndexOf(term, StringComparison.OrdinalIgnoreCase) >= 0 Then
                    score += 10 ' 文件名匹配得分较高
                End If
            Next

            ' 内容匹配
            If Not String.IsNullOrEmpty(fileInfo.Content) Then
                ' 找出匹配的行
                Dim lines = fileInfo.Content.Split({Environment.NewLine}, StringSplitOptions.RemoveEmptyEntries)
                For Each line In lines
                    Dim containsTerms As Boolean = False
                    For Each term In queryTerms
                        If line.IndexOf(term, StringComparison.OrdinalIgnoreCase) >= 0 Then
                            containsTerms = True

                            ' 如果是OCR文本，得分略低（降低误报带来的噪声）
                            If line.StartsWith("[图片OCR文本]") Then
                                score += 0.8
                            Else
                                score += 1
                            End If
                        End If
                    Next

                    If containsTerms Then
                        matchedLines.Add(line.Trim())
                    End If
                Next
            End If

            ' 关键词匹配
            If fileInfo.Keywords IsNot Nothing Then
                For Each term In queryTerms
                    If fileInfo.Keywords.Any(Function(k) k.IndexOf(term, StringComparison.OrdinalIgnoreCase) >= 0) Then
                        score += 5 ' 关键词匹配得分较高
                    End If
                Next
            End If

            ' 如果有匹配，添加到结果
            If score > 0 Then
                ' 对于图片OCR结果，添加标记
                Dim isOcrResult = IsImageFile(fileInfo.Extension) AndAlso matchedLines.Count > 0

                results.Add(New SearchResult() With {
                .File = fileInfo,
                .score = score,
                .matchedLines = matchedLines,
                .isOcrResult = isOcrResult
            })
            End If
        Next

        ' 按得分排序
        Return results.OrderByDescending(Function(r) r.Score).ToList()
    End Function

    ' 判断是否为图片文件
    Private Function IsImageFile(extension As String) As Boolean
        Dim imageExtensions As String() = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".gif"}
        Return imageExtensions.Contains(extension.ToLower())
    End Function

    ' 预处理查询
    Private Function PreprocessQuery(query As String) As List(Of String)
        If String.IsNullOrWhiteSpace(query) Then
            Return New List(Of String)()
        End If

        ' 拆分查询词并移除停用词
        Dim stopWords As New List(Of String)({"a", "an", "the", "and", "or", "but", "is", "are", "was", "were", "be", "to", "of", "in", "that", "have", "it", "for"})

        ' 分词
        Dim terms As New List(Of String)()
        For Each term In Regex.Split(query.ToLower(), "\s+")
            If term.Length > 1 AndAlso Not stopWords.Contains(term) Then
                terms.Add(term)
            End If
        Next

        Return terms
    End Function

    ' 高亮显示匹配的文本
    Public Function HighlightMatches(text As String, query As String) As String
        If String.IsNullOrEmpty(text) OrElse String.IsNullOrEmpty(query) Then
            Return text
        End If

        Dim terms = PreprocessQuery(query)
        If terms.Count = 0 Then
            Return text
        End If

        Dim result = text
        For Each term In terms
            result = Regex.Replace(
                result,
                $"({Regex.Escape(term)})",
                "<mark>$1</mark>",
                RegexOptions.IgnoreCase)
        Next

        Return result
    End Function
End Class
